package com.taoge.aop;

import com.taoge.annotation.LeakyBucketLimit;
import com.taoge.exception.RateLimitException;
import com.taoge.model.LeakyBucket;
import com.taoge.properties.RateLimitProperties;
import org.springframework.aop.MethodBeforeAdvice;

import java.lang.reflect.Method;

/**
 * 漏桶限流切面
 * @author Mr.Tao
 */
public class LeakyBucketLimitAspect implements MethodBeforeAdvice {

    private final RateLimitProperties limitProperties;

    public LeakyBucketLimitAspect(RateLimitProperties limitProperties) {
        this.limitProperties = limitProperties;
    }

    @Override
    public void before(Method method, Object[] args, Object target) {
        // 获取LeakyBucketLimit注解，如果方法和类上同时出现，以方法上的为准
        Class<?> clazz = method.getDeclaringClass();
        LeakyBucketLimit leakyBucketLimit;
        leakyBucketLimit = method.getAnnotation(LeakyBucketLimit.class);
        if (leakyBucketLimit == null) {
            leakyBucketLimit = clazz.getAnnotation(LeakyBucketLimit.class);
            if (leakyBucketLimit == null) {
                return;
            }
        }

        // 根据用户IP（可选）和接口方法，构造key
        String key = limitProperties.getKey(leakyBucketLimit.value(), leakyBucketLimit.key(), clazz.getSimpleName(), method.getName(), limitProperties.getStrategyConfig().getLeakyBucketLimitStrategy());

        LeakyBucket leakyBucket = LeakyBucket.getLeakyBucket(key);
        if (leakyBucket == null) {
            LeakyBucket.putLeakyBucket(key, leakyBucketLimit.capacity(), leakyBucketLimit.leakRate());
            return;
        }
        if (leakyBucket.isFull()) {
            throw new RateLimitException("请求过于频繁，请稍后重试！");
        }
    }

}

